#
# example008.py
#
# Adaptive Speech watermarking in wavelet domain based on Logarithm
# Send only the exponents, make compilant with G.711 A-Law, where:
# s0000000wxyz`a -> s000wxyz, wxyz/10 = allows 15 levels
# s0000001wxyz`a -> s001wxyz, wxyz/10 = allows 15 levels
# --------------------------
# s000000Xwxyz`a -> s001wxyz, wxyz/10 = allows 31 levels
# 4 bits for the levels (15) + 1 bit for mask
#
# Recipient receives : message_received
#
# Copyright (C) 2012 Robert Buj Gelonch
# Copyright (C) 2012 David Megias Jimenez
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
__author__ = "Robert Buj Gelonch, and David Megias Jimenez"
__copyright__ = "Copyright 2012, Robert Buj Gelonch and David Megias Jimenez"
__credits__ = ["Robert Buj Gelonch", "David Megias Jimenez"]
__license__ = "GPL"
__version__ = "3"
__maintainer__ = "Robert Buj"
__email__ = "rbuj@uoc.edu"
__status__ = "Development"
__docformat__ = 'plaintext'

from matplotlib.ticker import MultipleLocator
from numpy import copysign
from numpy import fix
from numpy import floor
from numpy import fromstring
from numpy import int16
from numpy import log10
from numpy import sign
from numpy import uint8
from numpy import unpackbits
from numpy import zeros
from pylab import bar
from pylab import figure
from pylab import grid
from pylab import plot
from pylab import show
from pylab import stem
from pylab import subplot
from pylab import suptitle
from pylab import title
from pylab import xlabel
from pylab import ylabel
from pywt import dwt
from scipy import arange
from scipy.io.wavfile import read
from scipy.io.wavfile import write
from scipy.signal import daub

print "example008.py"
print

#--------------------------------------------------------------
# Parameters of the watermarking
#--------------------------------------------------------------
message_str = "hello!"
message_bytes = fromstring(message_str, dtype=uint8)
message_bits = unpackbits(message_bytes, axis=-1)
secret = message_bits # secret message in the watermark
numtaps = 10 # Daubechies coefficient
num_levels = 15.0 # number of levels in the log quantification
frame_size = 4 # number os sample for each frame

#--------------------------------------------------------------
# Reading an audio file
#--------------------------------------------------------------
(sample_rate, samples) = read("../../../media/OSR_uk_000_0020_8k.wav")

#--------------------------------------------------------------
# Getting some time & freq data from audio file
#--------------------------------------------------------------
n = len(samples)
Ts = 1.0 / sample_rate
T = n / sample_rate
t = arange(0, n * Ts, Ts) # time vector

#--------------------------------------------------------------
# DWT - Daubechies wavelet transform
#--------------------------------------------------------------
daub_coeff = daub(numtaps)
cA, cD = dwt(samples, 'db10')

#--------------------------------------------------------------
# Marking the filtered message using logarithmic quantification
#--------------------------------------------------------------
max_value = max(abs(samples))
steps = max_value ** (1.0 / num_levels)
num_frames = floor(len(cA) / frame_size)
marked_message = zeros(len(cA))
idx = 0
for num_frame in range(int(num_frames)):
    if secret[idx] == 0:
        marked_message[num_frame * frame_size:(num_frame + 1) * frame_size] = \
            floor(log10(abs(cA[(num_frame * frame_size):(num_frame + 1) * \
                  frame_size])) / log10(steps)) * 4
    else:
        marked_message[num_frame * frame_size:(num_frame + 1) * frame_size] = \
            (floor(log10(abs(cA[(num_frame * frame_size):(num_frame + 1) * \
             frame_size])) / log10(steps)) + 0.5) * 4
    marked_message[num_frame * frame_size:(num_frame + 1) * frame_size] = \
        copysign(marked_message[num_frame * frame_size:(num_frame + 1) * frame_size], \
                 sign(cA[(num_frame * frame_size):(num_frame + 1) * frame_size]))
    if idx != len(secret)-1:
        idx += 1
    else:
        idx = 0
if num_frames != len(cA):
    print 'ToDo: fix last samples'

#--------------------------------------------------------------
# After G.711 A-LAW
#--------------------------------------------------------------
message_received = marked_message / 2

#--------------------------------------------------------------
# Getting the secret using the received message
#--------------------------------------------------------------
message_secret = message_received / 2
message_extracted = fix(message_secret)
message_secret_val = (abs(message_secret) - \
                      abs(message_extracted)) > 0.4

#--------------------------------------------------------------
# Restore audio from the received message
#--------------------------------------------------------------
restored_audio = copysign(steps ** abs(message_extracted), \
                          sign(message_received))

#--------------------------------------------------------------
# Plotting
#--------------------------------------------------------------
fig = figure(num=None, \
             figsize=(12, 9), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
suptitle('Adaptive Speech watermarking in wavelet domain based on Logarithm')

ax = subplot(211)
ylabel('Amplitude')
title('Filtered Signal: cA (db10)')
plot(cA, 'b')
ax.set_xticklabels([])
ax.set_xlim([0, len(cA)])
grid(True)

ax = subplot(212)
ylabel('Amplitude')
xlabel('Samples')
title('Marked cA')
plot(marked_message, 'b')
plot(message_received, 'r', label="After G.711 A-LAW")
ax.set_xlim([0, len(marked_message)])
grid(True)
leg = ax.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

# Save plot
plt.savefig("../plots/example008.py.1.png", format="png")

# Show plot
show()

#--------------------------------------------------------------
# Plotting
#--------------------------------------------------------------
fig = figure(num=None, \
             figsize=(14, 9), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
suptitle('Adaptive Speech watermarking in wavelet domain based on Logarithm')

ax = subplot(311)
ax.xaxis.set_major_locator(MultipleLocator(frame_size))
ax.set_xticklabels([])
ylabel('Amplitude')
title('Marked Frames')
grid(True)
stem(range(frame_size * len(secret)), \
     marked_message[:(frame_size * len(secret))], \
     'b')

ax = subplot(312)
ax.xaxis.set_major_locator(MultipleLocator(frame_size))
ax.set_xticklabels([])
title('Steganography')
ylabel('Base')
grid(True, which='both')
bar(range(frame_size * len(secret)), \
    message_secret[:(frame_size * len(secret))], \
    1, \
    color='r', \
    label="decimals")
bar(range(frame_size * len(secret)), \
    fix(message_secret[:(frame_size * len(secret))]), \
    1, \
    color='b', \
    label="floor")
leg = ax.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

ax = subplot(313)
title('Secret Message')
ylabel('Value')
xlabel('Frames')
ax.set_ylim([0, 1])
ax.set_autoscaley_on(False)
ax.set_yticks([0, 1])
ax.xaxis.set_major_locator(MultipleLocator(frame_size))
ax.set_xticklabels(range((frame_size * len(secret))))
grid(True)
plot(range(frame_size * len(secret)), \
     message_secret_val[:(frame_size * len(secret))], \
     'bo--')

# Save plot
plt.savefig("../plots/example008.py.2.png", format="png")

# Show plot
show()

#--------------------------------------------------------------
# Plotting
#--------------------------------------------------------------
fig = figure(num=None, \
             figsize=(14, 9), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
suptitle('Adaptive Speech watermarking in wavelet domain based on Logarithm')

ax = subplot(111)
ylabel('Amplitude')
title('Restored vs Filtered Signal')
plot(cA, 'b', label="cA")
plot(restored_audio, 'r', label="Restored")
ax.set_xticklabels([])
ax.set_xlim([0, len(cA)])
grid(True)
leg = ax.legend(loc='best', fancybox=True)
leg.get_frame().set_alpha(0.5)

# Save plot
plt.savefig("../plots/example008.py.3.png", format="png")

# Show plot
show()


#--------------------------------------------------------------
# Save to a file the rebuilded audio
#--------------------------------------------------------------
write("../example008.restored.wav", \
      sample_rate / 2, \
      restored_audio.astype(int16))

print "Done"